package com.hulman.oms.service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

@Service
@Slf4j
public class SSEService
{
    @Autowired
    private OllamaChatModel chatModel;

    public SseEmitter emitter(String message)
    {
        SseEmitter emitter = new SseEmitter(60_000L);
        Prompt prompt = new Prompt(new UserMessage(message));
        chatModel.stream(prompt).subscribe(response ->
        {
            try
            {
                AssistantMessage assistantMessage = response.getResult().getOutput();
                String text = assistantMessage.getText();
                //log.info("text:{}", text);
                emitter.send(SseEmitter.event().data(text).id(String.valueOf(System.currentTimeMillis())).build());
            } catch (Exception e)
            {
                emitter.completeWithError(e);
            }
        }, error ->
        {
            emitter.completeWithError(error);
        }, () ->
        {
            try
            {
                emitter.send(SseEmitter.event().data("END").id(String.valueOf(System.currentTimeMillis())).build());
            } catch (IOException e)
            {
                throw new RuntimeException(e);
            }
            emitter.complete();
        });
        return emitter;
    }
}
